import os

from dotenv import load_dotenv
from langchain_community.agent_toolkits import SQLDatabaseToolkit
from langchain_community.utilities import SQLDatabase
from langchain_core.messages import SystemMessage, HumanMessage
from langchain_core.prompts import ChatPromptTemplate
from langchain_openai import ChatOpenAI
from langgraph.prebuilt import chat_agent_executor

load_dotenv()

# 1.创建模型
model = ChatOpenAI(
    model='qwen-plus',
    api_key=os.getenv("DASHSCOPE_API_KEY"),
    base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",
)

# SQLAIChemy 数据库工具

HOST_NAME = '127.0.0.1'
PORT = '3306'
# DATABASE = 'ruoyi-vue-pro'
DATABASE = 'wx'
USERNAME = 'root'
PASSWORD = 'root123456'
# mysqlclient 驱动URL
MYSQL_URI = 'mysql+mysqldb://{}:{}@{}:{}/{}?charset=utf8mb4'.format(USERNAME, PASSWORD, HOST_NAME, PORT, DATABASE)

db = SQLDatabase.from_uri(MYSQL_URI)

# SQL工具
toolkit = SQLDatabaseToolkit(db=db, llm=model)
tools = toolkit.get_tools()


# 使用agent完整整个数据库的整合
system_prompt = """
您是一个被设计用来与SQL数据库交互的代理。
给定一个输入问题，创建一个语法正确的SQL语句并执行，然后查看查询结果并返回答案。
除非用户指定了他们想要法得的示例的具体数量，否则始终将SQL查询限制为最多10个结果。
你可以按相关列对结果进行排序，以返回MySQL数据库中最匹配的数据。
您可以使用与数据库交互的工具。在执行查询之前，你必须仔细检查。如果在执行查询时出现错误，请重写查询SQL并重试。
不要对数据库做任何DML语句(插入，更新，删除，删除等)。

首先，你应该查看数据库中的表，看看可以查询什么。
不要跳过这一步。
然后查询最相关的表的模式。
"""
system_message = SystemMessage(content=system_prompt)

# # 1. 创建带系统消息的提示模板
# prompt = ChatPromptTemplate.from_messages([
#     ("system", system_prompt),  # 系统提示词
# ])


# 创建agent
# agent_executor = chat_agent_executor.create_tool_calling_executor(model=model, tools=tools, system_message) # 会报错
agent_executor = chat_agent_executor.create_react_agent(model=model, tools=tools, prompt=system_message)
# agent_executor = chat_agent_executor.create_tool_calling_executor(model=model, tools=tools, prompt=prompt)
# agent_executor = chat_agent_executor.create_tool_calling_executor(model=model, tools=tools)


# create_react_agent 和 create_tool_calling_executor 这两个函数有什么区别？？？


resp = agent_executor.invoke({'messages': [HumanMessage(content='请问：菜单表中有多少条数据？')]})
# resp = agent_executor.invoke({'messages': [HumanMessage(content='请问：菜单中有多少种角色？')]})

result = resp['messages']
print(result)
print(len(result))
print(result[len(result)-1])

